package com.example.myspringai.util;

import org.apache.commons.lang3.StringUtils;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.analysis.ar.ArabicAnalyzer;
import org.apache.lucene.analysis.tokenattributes.CharTermAttribute;
import org.apache.poi.ss.usermodel.Row;
import org.apache.poi.ss.usermodel.Sheet;
import org.apache.poi.ss.usermodel.Workbook;
import org.apache.poi.xssf.usermodel.XSSFWorkbook;

import java.io.FileInputStream;
import java.io.IOException;
import java.io.StringReader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class ArabicWER {

    private static final String[] PUNCTUATIONS = {".", ",", "،", ":", "؛", "؟", "!", "(", ")", "[", "]", "{", "}", "\"", "'"};

    public static List<String> tokenize(String text) throws IOException {
        // Remove punctuation
        for (String punctuation : PUNCTUATIONS) {
            text = text.replace(punctuation, "");
        }

        // Initialize the Arabic analyzer
        ArabicAnalyzer analyzer = new ArabicAnalyzer();
        List<String> result = new ArrayList<>();

        // Use tokenStream to tokenize the text
        TokenStream tokenStream = analyzer.tokenStream(null, new StringReader(text));
        CharTermAttribute attr = tokenStream.addAttribute(CharTermAttribute.class);
        tokenStream.reset();

        // Collect tokens
        while (tokenStream.incrementToken()) {
            result.add(attr.toString());
        }

        tokenStream.end();
        tokenStream.close();

        return result;
    }

    public static double calculateWER(List<String> ref, List<String> hyp) {
        int[] costs = new int[ref.size() + 1];
        for (int i = 0; i <= ref.size(); i++) {
            costs[i] = i;
        }

        for (int j = 1; j <= hyp.size(); j++) {
            int lastValue = costs[0];
            costs[0] = j;
            for (int i = 1; i <= ref.size(); i++) {
                int newValue = Math.min(1 + Math.min(costs[i], costs[i - 1]), ref.get(i - 1).equals(hyp.get(j - 1)) ? lastValue : lastValue + 1);
                lastValue = costs[i];
                costs[i] = newValue;
            }
        }

        double result =  (double) costs[ref.size()] / ref.size()  ;
        //此处保留两位小数，四舍五入
        return Math.round(result * 100) / 100.0;
    }

    public static void main(String[] args) {
        excelTest();
    }

    public static void excelTest(){
        String filePath = "D:\\语音转写厂商备选清单.xlsx";
        // 创建一个List数据
        List<List<String>> data =new ArrayList<>( Arrays.asList(
                Arrays.asList("数据集标注","whisper-GUI","whisper-small-egyptian-arabic-All","intella.me","中译语通"))
        );
        try (FileInputStream fis = new FileInputStream(filePath); Workbook workbook = new XSSFWorkbook(fis)) {
            Sheet sheet = workbook.getSheetAt(2);
            for (Row row : sheet) {
                if (row.getRowNum() == 0 || row.getRowNum() == 1) {
                    continue; // Skip the header row
                }
                if (row.getRowNum() > 100) {
                    continue; // Skip the header row
                }

                String annotateText = row.getCell(2) == null ? "" : String.valueOf(row.getCell(2).getStringCellValue());
                String whisperGuiText = row.getCell(3) == null ? "" : String.valueOf(row.getCell(3).getStringCellValue());
                String whisperSmallText = row.getCell(4) == null ? "" : String.valueOf(row.getCell(4).getStringCellValue());
                String intellaText = row.getCell(5) == null ? "" : String.valueOf(row.getCell(5).getStringCellValue());
                String zhongYiText = row.getCell(6) == null ? "" : String.valueOf(row.getCell(6).getStringCellValue());
                if(StringUtils.isEmpty(annotateText)){
                    continue;
                }
                List<String> rowData = new ArrayList<>();
                List<String> refTokens = tokenize(annotateText);
                rowData.add(annotateText);
                rowData.add(calculateWER(refTokens, tokenize(whisperGuiText))+"");
                rowData.add(calculateWER(refTokens, tokenize(whisperSmallText))+"");
                rowData.add(calculateWER(refTokens, tokenize(intellaText))+"");
                rowData.add(calculateWER(refTokens, tokenize(zhongYiText))+"");
                System.out.println(row.getRowNum()+":"+rowData);
                data.add(rowData);
            }
            ExcelUtil.writeExcel(data, "stt-mix");
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}
